// Copyright 2020 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/browser/ash/cert_provisioning/cert_provisioning_worker.h"

#include <memory>
#include <string>

#include "base/base64.h"
#include "base/callback.h"
#include "base/json/json_string_value_serializer.h"
#include "base/json/json_writer.h"
#include "base/strings/stringprintf.h"
#include "base/test/gmock_callback_support.h"
#include "base/test/metrics/histogram_tester.h"
#include "base/test/values_test_util.h"
#include "base/time/time.h"
#include "chrome/browser/ash/attestation/mock_tpm_challenge_key_subtle.h"
#include "chrome/browser/ash/attestation/tpm_challenge_key_subtle.h"
#include "chrome/browser/ash/cert_provisioning/cert_provisioning_common.h"
#include "chrome/browser/ash/cert_provisioning/cert_provisioning_metrics.h"
#include "chrome/browser/ash/cert_provisioning/cert_provisioning_test_helpers.h"
#include "chrome/browser/ash/cert_provisioning/mock_cert_provisioning_invalidator.h"
#include "chrome/browser/ash/platform_keys/key_permissions/fake_user_private_token_kpm_service.h"
#include "chrome/browser/ash/platform_keys/key_permissions/key_permissions_manager.h"
#include "chrome/browser/ash/platform_keys/key_permissions/key_permissions_manager_impl.h"
#include "chrome/browser/ash/platform_keys/key_permissions/mock_key_permissions_manager.h"
#include "chrome/browser/ash/platform_keys/key_permissions/user_private_token_kpm_service_factory.h"
#include "chrome/browser/ash/platform_keys/mock_platform_keys_service.h"
#include "chrome/browser/ash/platform_keys/platform_keys_service.h"
#include "chrome/browser/ash/platform_keys/platform_keys_service_factory.h"
#include "chrome/browser/platform_keys/platform_keys.h"
#include "chromeos/dbus/attestation/fake_attestation_client.h"
#include "components/policy/core/common/cloud/mock_cloud_policy_client.h"
#include "components/prefs/pref_change_registrar.h"
#include "components/prefs/pref_service.h"
#include "components/prefs/testing_pref_service.h"
#include "content/public/test/browser_task_environment.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace ash {
namespace cert_provisioning {
namespace {

namespace em = ::enterprise_management;

using ::ash::attestation::MockTpmChallengeKeySubtle;
using ::ash::platform_keys::KeyUsage;
using ::base::test::IsJson;
using ::base::test::ParseJson;
using ::base::test::RunOnceCallback;
using ::chromeos::platform_keys::HashAlgorithm;
using ::chromeos::platform_keys::KeyAttributeType;
using ::chromeos::platform_keys::Status;
using ::chromeos::platform_keys::TokenId;
using ::testing::_;
using ::testing::AtLeast;
using ::testing::Mock;
using ::testing::SaveArg;
using ::testing::StrictMock;

// Generated by chrome/test/data/policy/test_certs/create_test_certs.sh
constexpr char kFakeCertificate[] = R"(-----BEGIN CERTIFICATE-----
MIIDJzCCAg+gAwIBAgIBATANBgkqhkiG9w0BAQsFADAXMRUwEwYDVQQDDAxyb290
X2NhX2NlcnQwHhcNMjAwMjI1MTUyNTU2WhcNMzAwMjIyMTUyNTU2WjAUMRIwEAYD
VQQDDAkxMjcuMC4wLjEwggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDW
druvpaJovmyWzIcjtsSk/lp319+zNPSYGLzJzTeEmnFoDf/b89ft6xR1NIahmvVd
UHGOMlzgDKnNkqWw+pgpn6U8dk+leWnwlUefzDz7OY8qXfX29Vh0m/kATQc64lnp
rX19fEi2DOgH6heCQDSaHI/KAnAXccwl8kdGuTEnvdzbdHqQq8pPGpEqzC/NOjk7
kDNkUt0J74ZVMm4+jhVOgZ35mFLtC+xjfycBgbnt8yfPOzmOMwXTjYDPNaIy32AZ
t66oIToteoW5Ilg+j5Mto3unBDHrw8rml3+W/nwHuOPEIgBqLQFfWtXpuX8CbcS6
SFNK4hxCJOvlzUbgTpsrAgMBAAGjgYAwfjAMBgNVHRMBAf8EAjAAMB0GA1UdDgQW
BBRDEl1/2pL5LtKnpIly+XCj3N6MwDAfBgNVHSMEGDAWgBQrwVEnUQZlX850A2N+
URfS8BxoyzAdBgNVHSUEFjAUBggrBgEFBQcDAQYIKwYBBQUHAwIwDwYDVR0RBAgw
BocEfwAAATANBgkqhkiG9w0BAQsFAAOCAQEAXZd+Ul7GUFZPLSiTZ618hUI2UdO0
7rtPwBw3TephWuyEeHht+WhzA3sRL3nprEiJqIg5w/Tlfz4dsObpSU3vKmDhLzAx
HJrN5vKdbEj9wyuhYSRJwvJka1ZOgPzhQcDQOp1SqonNxLx/sSMDR2UIDMBGzrkQ
sDkn58N5eWm+hZADOAKROHR47j85VcsmYGK7z2x479YzsyWyOm0dbACXv7/HvFkz
56KvgxRaPZQzQUg5yuXa21IjQz07wyWSYnHpm2duAbYFl6CTR9Rlj5vpRkKsQP1W
mMhGDBfgEskdbM+0agsZrJupoQMBUbD5gflcJlW3kwlboi3dTtiGixfYWw==
-----END CERTIFICATE-----)";

// Extracted from the certificate using the command:
// openssl x509 -pubkey -noout -in cert.pem
// and reformatted as a single line.
constexpr char kPublicKeyBase64[] =
    "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA1na7r6WiaL5slsyHI7bEpP5ad9ffsz"
    "T0mBi8yc03hJpxaA3/2/"
    "PX7esUdTSGoZr1XVBxjjJc4AypzZKlsPqYKZ+lPHZPpXlp8JVHn8w8+"
    "zmPKl319vVYdJv5AE0HOuJZ6a19fXxItgzoB+"
    "oXgkA0mhyPygJwF3HMJfJHRrkxJ73c23R6kKvKTxqRKswvzTo5O5AzZFLdCe+"
    "GVTJuPo4VToGd+ZhS7QvsY38nAYG57fMnzzs5jjMF042AzzWiMt9gGbeuqCE6LXqFuSJYPo+"
    "TLaN7pwQx68PK5pd/lv58B7jjxCIAai0BX1rV6bl/Am3EukhTSuIcQiTr5c1G4E6bKwIDAQAB";

constexpr char kCertProfileId[] = "cert_profile_1";
constexpr char kCertProfileName[] = "Certificate Profile 1";
constexpr char kCertProfileVersion[] = "cert_profile_version_1";
constexpr base::TimeDelta kCertProfileRenewalPeriod = base::Seconds(0);
// Prefix + certificate profile name.
constexpr char kCertScopeStrUser[] = "google/chromeos/user";
constexpr char kCertScopeStrDevice[] = "google/chromeos/device";
constexpr char kInvalidationTopic[] = "fake_invalidation_topic_1";
constexpr char kDataToSign[] = "fake_data_to_sign_1";
constexpr char kChallenge[] = "fake_va_challenge_1";
constexpr char kChallengeResponse[] = "fake_va_challenge_response_1";
constexpr char kSignature[] = "fake_signature_1";
constexpr unsigned int kNonVaKeyModulusLengthBits = 2048;

const std::string& GetPublicKey() {
  static std::string public_key;
  if (public_key.empty()) {
    base::Base64Decode(kPublicKeyBase64, &public_key);
  }
  return public_key;
}

void VerifyDeleteKeyCalledOnce(CertScope cert_scope) {
  const std::vector<::attestation::DeleteKeysRequest> delete_keys_history =
      chromeos::AttestationClient::Get()
          ->GetTestInterface()
          ->delete_keys_history();
  EXPECT_EQ(delete_keys_history.size(), 1);
  EXPECT_EQ(delete_keys_history[0].username().empty(),
            cert_scope != CertScope::kUser);
  EXPECT_EQ(delete_keys_history[0].key_label_match(),
            GetKeyName(kCertProfileId));
  EXPECT_EQ(delete_keys_history[0].match_behavior(),
            ::attestation::DeleteKeysRequest::MATCH_BEHAVIOR_EXACT);
}

// Using macros to reduce boilerplate code, but keep real line numbers in
// error messages in case of expectation failure. They use some of protected
// fields of CertProvisioningWorkerTest class and may be considered as extra
// methods of it. *_OK macros immediately call callbacks with some successful
// results. *_NO_OP doesn't call callbacks.
#define EXPECT_PREPARE_KEY_OK(MOCK_TPM_CHALLENGE_KEY, PREPARE_KEY_FUNC)    \
  {                                                                        \
    auto public_key_result =                                               \
        attestation::TpmChallengeKeyResult::MakePublicKey(GetPublicKey()); \
    EXPECT_CALL((MOCK_TPM_CHALLENGE_KEY), PREPARE_KEY_FUNC)                \
        .Times(1)                                                          \
        .WillOnce(RunOnceCallback<4>(public_key_result));                  \
  }

#define EXPECT_SIGN_CHALLENGE_OK(MOCK_TPM_CHALLENGE_KEY, SIGN_CHALLENGE_FUNC) \
  {                                                                           \
    auto sign_challenge_result =                                              \
        attestation::TpmChallengeKeyResult::MakeChallengeResponse(            \
            kChallengeResponse);                                              \
    EXPECT_CALL((MOCK_TPM_CHALLENGE_KEY), SIGN_CHALLENGE_FUNC)                \
        .Times(1)                                                             \
        .WillOnce(RunOnceCallback<1>(sign_challenge_result));                 \
  }

#define EXPECT_REGISTER_KEY_OK(MOCK_TPM_CHALLENGE_KEY, REGISTER_KEY_FUNC) \
  {                                                                       \
    auto register_key_result =                                            \
        attestation::TpmChallengeKeyResult::MakeSuccess();                \
    EXPECT_CALL((MOCK_TPM_CHALLENGE_KEY), REGISTER_KEY_FUNC)              \
        .Times(1)                                                         \
        .WillOnce(RunOnceCallback<0>(register_key_result));               \
  }

#define EXPECT_START_CSR_OK(START_CSR_FUNC, HASHING_ALGO)             \
  {                                                                   \
    EXPECT_CALL(cloud_policy_client_, START_CSR_FUNC)                 \
        .Times(1)                                                     \
        .WillOnce(RunOnceCallback<4>(                                 \
            policy::DeviceManagementStatus::DM_STATUS_SUCCESS,        \
            /*response_error=*/absl::nullopt,                         \
            /*try_again_later_ms=*/absl::nullopt, kInvalidationTopic, \
            kChallenge, HASHING_ALGO, kDataToSign));                  \
  }

#define EXPECT_START_CSR_OK_WITHOUT_VA(START_CSR_FUNC, HASHING_ALGO)  \
  {                                                                   \
    EXPECT_CALL(cloud_policy_client_, START_CSR_FUNC)                 \
        .Times(1)                                                     \
        .WillOnce(RunOnceCallback<4>(                                 \
            policy::DeviceManagementStatus::DM_STATUS_SUCCESS,        \
            /*response_error=*/absl::nullopt,                         \
            /*try_again_later_ms=*/absl::nullopt, kInvalidationTopic, \
            /*va_challenge=*/"", HASHING_ALGO, kDataToSign));         \
  }

#define EXPECT_START_CSR_TRY_LATER(START_CSR_FUNC, DELAY_MS)       \
  {                                                                \
    EXPECT_CALL(cloud_policy_client_, START_CSR_FUNC)              \
        .Times(1)                                                  \
        .WillOnce(RunOnceCallback<4>(                              \
            policy::DeviceManagementStatus::DM_STATUS_SUCCESS,     \
            /*response_error=*/absl::nullopt,                      \
            /*try_again_later_ms=*/(DELAY_MS), kInvalidationTopic, \
            /*va_challenge=*/"",                                   \
            enterprise_management::HashingAlgorithm::              \
                HASHING_ALGORITHM_UNSPECIFIED,                     \
            /*data_to_sign=*/""));                                 \
  }

#define EXPECT_START_CSR_INVALID_REQUEST(START_CSR_FUNC)                     \
  {                                                                          \
    EXPECT_CALL(cloud_policy_client_, START_CSR_FUNC)                        \
        .Times(1)                                                            \
        .WillOnce(RunOnceCallback<4>(                                        \
            policy::DeviceManagementStatus::DM_STATUS_REQUEST_INVALID,       \
            /*response_error=*/absl::nullopt,                                \
            /*try_again_later_ms=*/absl::nullopt, /*invalidation_topic=*/"", \
            /*va_challenge=*/"",                                             \
            enterprise_management::HashingAlgorithm::                        \
                HASHING_ALGORITHM_UNSPECIFIED,                               \
            /*data_to_sign=*/""));                                           \
  }

#define EXPECT_START_CSR_CA_ERROR(START_CSR_FUNC)                            \
  {                                                                          \
    EXPECT_CALL(cloud_policy_client_, START_CSR_FUNC)                        \
        .Times(1)                                                            \
        .WillOnce(RunOnceCallback<4>(                                        \
            policy::DeviceManagementStatus::DM_STATUS_SUCCESS,               \
            /*response_error=*/CertProvisioningResponseError::CA_ERROR,      \
            /*try_again_later_ms=*/absl::nullopt, /*invalidation_topic=*/"", \
            /*va_challenge=*/"",                                             \
            enterprise_management::HashingAlgorithm::                        \
                HASHING_ALGORITHM_UNSPECIFIED,                               \
            /*data_to_sign=*/""));                                           \
  }

#define EXPECT_START_CSR_TEMPORARY_UNAVAILABLE(START_CSR_FUNC)               \
  {                                                                          \
    EXPECT_CALL(cloud_policy_client_, START_CSR_FUNC)                        \
        .Times(1)                                                            \
        .WillOnce(RunOnceCallback<4>(                                        \
            policy::DeviceManagementStatus::DM_STATUS_TEMPORARY_UNAVAILABLE, \
            /*response_error=*/absl::nullopt,                                \
            /*try_again_later_ms=*/absl::nullopt, /*invalidation_topic=*/"", \
            /*va_challenge=*/"",                                             \
            enterprise_management::HashingAlgorithm::                        \
                HASHING_ALGORITHM_UNSPECIFIED,                               \
            /*data_to_sign=*/""));                                           \
  }

#define EXPECT_START_CSR_SERVICE_ACTIVATION_PENDING(START_CSR_FUNC)            \
  {                                                                            \
    EXPECT_CALL(cloud_policy_client_, START_CSR_FUNC)                          \
        .Times(1)                                                              \
        .WillOnce(RunOnceCallback<4>(policy::DeviceManagementStatus::          \
                                         DM_STATUS_SERVICE_ACTIVATION_PENDING, \
                                     /*response_error=*/absl::nullopt,         \
                                     /*try_again_later_ms=*/absl::nullopt,     \
                                     /*invalidation_topic=*/"",                \
                                     /*va_challenge=*/"",                      \
                                     enterprise_management::HashingAlgorithm:: \
                                         HASHING_ALGORITHM_UNSPECIFIED,        \
                                     /*data_to_sign=*/""));                    \
  }

#define EXPECT_START_CSR_INCONSISTENT_DATA(START_CSR_FUNC)                   \
  {                                                                          \
    EXPECT_CALL(cloud_policy_client_, START_CSR_FUNC)                        \
        .Times(1)                                                            \
        .WillOnce(RunOnceCallback<4>(                                        \
            policy::DeviceManagementStatus::                                 \
                DM_STATUS_SUCCESS, /*response_error=*/                       \
            CertProvisioningResponseError::INCONSISTENT_DATA,                \
            /*try_again_later_ms=*/absl::nullopt, /*invalidation_topic=*/"", \
            /*va_challenge=*/"",                                             \
            enterprise_management::HashingAlgorithm::                        \
                HASHING_ALGORITHM_UNSPECIFIED,                               \
            /*data_to_sign=*/""));                                           \
  }

#define EXPECT_START_CSR_NO_OP(START_CSR_FUNC) \
  { EXPECT_CALL(cloud_policy_client_, START_CSR_FUNC).Times(1); }

#define EXPECT_FINISH_CSR_OK(FINISH_CSR_FUNC)                                 \
  {                                                                           \
    EXPECT_CALL(cloud_policy_client_, FINISH_CSR_FUNC)                        \
        .Times(1)                                                             \
        .WillOnce(RunOnceCallback<6>(                                         \
            policy::DeviceManagementStatus::DM_STATUS_SUCCESS, absl::nullopt, \
            absl::nullopt));                                                  \
  }

#define EXPECT_FINISH_CSR_TRY_LATER(FINISH_CSR_FUNC, DELAY_MS)                \
  {                                                                           \
    EXPECT_CALL(cloud_policy_client_, FINISH_CSR_FUNC)                        \
        .Times(1)                                                             \
        .WillOnce(RunOnceCallback<6>(                                         \
            policy::DeviceManagementStatus::DM_STATUS_SUCCESS, absl::nullopt, \
            /*try_again_later_ms=*/(DELAY_MS)));                              \
  }

#define EXPECT_FINISH_CSR_SERVICE_ACTIVATION_PENDING(FINISH_CSR_FUNC)          \
  {                                                                            \
    EXPECT_CALL(cloud_policy_client_, FINISH_CSR_FUNC)                         \
        .Times(1)                                                              \
        .WillOnce(RunOnceCallback<6>(policy::DeviceManagementStatus::          \
                                         DM_STATUS_SERVICE_ACTIVATION_PENDING, \
                                     absl::nullopt, absl::nullopt));           \
  }

#define EXPECT_DOWNLOAD_CERT_OK(DOWNLOAD_CERT_FUNC)                           \
  {                                                                           \
    EXPECT_CALL(cloud_policy_client_, DOWNLOAD_CERT_FUNC)                     \
        .Times(1)                                                             \
        .WillOnce(RunOnceCallback<4>(                                         \
            policy::DeviceManagementStatus::DM_STATUS_SUCCESS, absl::nullopt, \
            absl::nullopt, kFakeCertificate));                                \
  }

#define EXPECT_DOWNLOAD_CERT_SERVICE_ACTIVATION_PENDING(DOWNLOAD_CERT_FUNC)    \
  {                                                                            \
    EXPECT_CALL(cloud_policy_client_, DOWNLOAD_CERT_FUNC)                      \
        .Times(1)                                                              \
        .WillOnce(RunOnceCallback<4>(policy::DeviceManagementStatus::          \
                                         DM_STATUS_SERVICE_ACTIVATION_PENDING, \
                                     absl::nullopt, absl::nullopt,             \
                                     kFakeCertificate));                       \
  }

#define EXPECT_DOWNLOAD_CERT_TRY_LATER(DOWNLOAD_CERT_FUNC, DELAY_MS)          \
  {                                                                           \
    EXPECT_CALL(cloud_policy_client_, DOWNLOAD_CERT_FUNC)                     \
        .Times(1)                                                             \
        .WillOnce(RunOnceCallback<4>(                                         \
            policy::DeviceManagementStatus::DM_STATUS_SUCCESS, absl::nullopt, \
            /*try_again_later_ms=*/(DELAY_MS), /*certificate=*/""));          \
  }

#define EXPECT_DOWNLOAD_CERT_NO_OP(DOWNLOAD_CERT_FUNC) \
  { EXPECT_CALL(cloud_policy_client_, DOWNLOAD_CERT_FUNC).Times(1); }

#define EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SET_FUNC)        \
  {                                                      \
    EXPECT_CALL(*platform_keys_service_, SET_FUNC)       \
        .Times(1)                                        \
        .WillOnce(RunOnceCallback<4>(Status::kSuccess)); \
  }

#define EXPECT_SET_ATTRIBUTE_FOR_KEY_FAIL(SET_FUNC)            \
  {                                                            \
    EXPECT_CALL(*platform_keys_service_, SET_FUNC)             \
        .Times(1)                                              \
        .WillOnce(RunOnceCallback<4>(Status::kErrorInternal)); \
  }

#define EXPECT_SIGN_RSAPKC1_DIGEST_OK(SIGN_FUNC)                     \
  {                                                                  \
    EXPECT_CALL(*platform_keys_service_, SIGN_FUNC)                  \
        .Times(1)                                                    \
        .WillOnce(RunOnceCallback<4>(kSignature, Status::kSuccess)); \
  }

#define EXPECT_SIGN_RSAPKC1_RAW_OK(SIGN_FUNC)                        \
  {                                                                  \
    EXPECT_CALL(*platform_keys_service_, SIGN_FUNC)                  \
        .Times(1)                                                    \
        .WillOnce(RunOnceCallback<3>(kSignature, Status::kSuccess)); \
  }

#define EXPECT_SIGN_RSAPKC1_DIGEST_FAIL(SIGN_FUNC)                         \
  {                                                                        \
    EXPECT_CALL(*platform_keys_service_, SIGN_FUNC)                        \
        .Times(1)                                                          \
        .WillOnce(                                                         \
            RunOnceCallback<4>(/*signature=*/"", Status::kErrorInternal)); \
  }

#define EXPECT_IMPORT_CERTIFICATE_OK(IMPORT_FUNC)        \
  {                                                      \
    EXPECT_CALL(*platform_keys_service_, IMPORT_FUNC)    \
        .Times(1)                                        \
        .WillOnce(RunOnceCallback<2>(Status::kSuccess)); \
  }

// A mock for observing the result callback of the worker.
class CallbackObserver {
 public:
  MOCK_METHOD(void,
              Callback,
              (const CertProfile& profile, CertProvisioningWorkerState state));
};

// A mock for observing the state change callback of the worker.
class StateChangeCallbackObserver {
 public:
  MOCK_METHOD(void, StateChangeCallback, ());
};

class CertProvisioningWorkerTest : public ::testing::Test {
 public:
  CertProvisioningWorkerTest() { Init(); }
  CertProvisioningWorkerTest(const CertProvisioningWorkerTest&) = delete;
  CertProvisioningWorkerTest& operator=(const CertProvisioningWorkerTest&) =
      delete;
  ~CertProvisioningWorkerTest() override = default;

  void SetUp() override {
    ::chromeos::AttestationClient::InitializeFake();
    // There should not be any calls to callback before this expect is
    // overridden.
    EXPECT_CALL(callback_observer_, Callback).Times(0);

    RegisterProfilePrefs(testing_pref_service_.registry());
    RegisterLocalStatePrefs(testing_pref_service_.registry());
  }

  void TearDown() override {
    EXPECT_FALSE(
        attestation::TpmChallengeKeySubtleFactory::WillReturnTestingInstance());
    ::chromeos::AttestationClient::Shutdown();
  }

 protected:
  void Init() {
    platform_keys_service_ =
        static_cast<platform_keys::MockPlatformKeysService*>(
            platform_keys::PlatformKeysServiceFactory::GetInstance()
                ->SetTestingFactoryAndUse(
                    GetProfile(),
                    base::BindRepeating(
                        &platform_keys::BuildMockPlatformKeysService)));
    ASSERT_TRUE(platform_keys_service_);
    platform_keys::PlatformKeysServiceFactory::GetInstance()
        ->SetDeviceWideServiceForTesting(platform_keys_service_);

    key_permissions_manager_ =
        std::make_unique<platform_keys::MockKeyPermissionsManager>();

    platform_keys::UserPrivateTokenKeyPermissionsManagerServiceFactory::
        GetInstance()
            ->SetTestingFactory(
                GetProfile(),
                base::BindRepeating(
                    &platform_keys::
                        BuildFakeUserPrivateTokenKeyPermissionsManagerService,
                    key_permissions_manager_.get()));
    platform_keys::KeyPermissionsManagerImpl::
        SetSystemTokenKeyPermissionsManagerForTesting(
            key_permissions_manager_.get());

    // Only explicitly expected removals are allowed.
    EXPECT_CALL(*platform_keys_service_, RemoveCertificate).Times(0);
    EXPECT_CALL(*platform_keys_service_, RemoveKey).Times(0);
  }

  void FastForwardBy(base::TimeDelta delta) {
    task_environment_.FastForwardBy(delta);
  }

  // Replaces next result of TpmChallengeKeySubtleFactory and return pointer to
  // the mock. The mock will injected into next created worker and will live
  // until worker's destruction. Should be called before creation of every
  // worker.
  MockTpmChallengeKeySubtle* PrepareTpmChallengeKey() {
    auto mock_tpm_challenge_key_subtle_impl =
        std::make_unique<MockTpmChallengeKeySubtle>();

    MockTpmChallengeKeySubtle* tpm_challenge_key_impl =
        mock_tpm_challenge_key_subtle_impl.get();

    attestation::TpmChallengeKeySubtleFactory::SetForTesting(
        std::move(mock_tpm_challenge_key_subtle_impl));

    CHECK(tpm_challenge_key_impl);
    return tpm_challenge_key_impl;
  }

  base::RepeatingClosure GetStateChangeCallback() {
    return base::BindRepeating(
        &StateChangeCallbackObserver ::StateChangeCallback,
        base::Unretained(&state_change_callback_observer_));
  }

  CertProvisioningWorkerCallback GetResultCallback() {
    return base::BindOnce(&CallbackObserver::Callback,
                          base::Unretained(&callback_observer_));
  }

  Profile* GetProfile() { return profile_helper_for_testing_.GetProfile(); }

  std::unique_ptr<MockCertProvisioningInvalidator> MakeInvalidator() {
    return std::make_unique<MockCertProvisioningInvalidator>();
  }

  std::unique_ptr<MockCertProvisioningInvalidator> MakeInvalidator(
      MockCertProvisioningInvalidator** mock_invalidator) {
    auto result = std::make_unique<MockCertProvisioningInvalidator>();
    *mock_invalidator = result.get();
    return result;
  }

  content::BrowserTaskEnvironment task_environment_{
      base::test::TaskEnvironment::TimeSource::MOCK_TIME};

  StrictMock<StateChangeCallbackObserver> state_change_callback_observer_;
  StrictMock<CallbackObserver> callback_observer_;
  ProfileHelperForTesting profile_helper_for_testing_;
  TestingPrefServiceSimple testing_pref_service_;

  policy::MockCloudPolicyClient cloud_policy_client_;
  platform_keys::MockPlatformKeysService* platform_keys_service_ = nullptr;
  std::unique_ptr<platform_keys::MockKeyPermissionsManager>
      key_permissions_manager_;
};

// Checks that the worker makes all necessary requests to other modules during
// success scenario.
TEST_F(CertProvisioningWorkerTest, Success) {
  base::HistogramTester histogram_tester;

  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  MockCertProvisioningInvalidator* mock_invalidator = nullptr;
  CertProvisioningWorkerImpl worker(
      CertScope::kUser, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(&mock_invalidator),
      GetStateChangeCallback(), GetResultCallback());

  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_START_CSR_OK(
        ClientCertProvisioningStartCsr(kCertScopeStrUser, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        em::HashingAlgorithm::SHA256);
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_CALL(*mock_invalidator, Register(kInvalidationTopic, _)).Times(1);

    EXPECT_SIGN_CHALLENGE_OK(*mock_tpm_challenge_key,
                             StartSignChallengeStep(kChallenge,
                                                    /*callback=*/_));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_REGISTER_KEY_OK(*mock_tpm_challenge_key, StartRegisterKeyStep);
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_CALL(
        *key_permissions_manager_,
        AllowKeyForUsage(/*callback=*/_, KeyUsage::kCorporate, GetPublicKey()));

    EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
        TokenId::kUser, GetPublicKey(),
        KeyAttributeType::kCertificateProvisioningId, kCertProfileId, _));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_SIGN_RSAPKC1_DIGEST_OK(
        SignRSAPKCS1Digest(::testing::Optional(TokenId::kUser), kDataToSign,
                           GetPublicKey(), HashAlgorithm::HASH_ALGORITHM_SHA256,
                           /*callback=*/_));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_FINISH_CSR_OK(ClientCertProvisioningFinishCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        kChallengeResponse, kSignature, /*callback=*/_));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_DOWNLOAD_CERT_OK(ClientCertProvisioningDownloadCert(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));

    EXPECT_IMPORT_CERTIFICATE_OK(ImportCertificate(TokenId::kUser,
                                                   /*certificate=*/_,
                                                   /*callback=*/_));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_CALL(*mock_invalidator, Unregister()).Times(1);

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kSucceeded))
        .Times(1);
  }

  worker.DoStep();

  histogram_tester.ExpectUniqueSample("ChromeOS.CertProvisioning.Result.User",
                                      CertProvisioningWorkerState::kSucceeded,
                                      1);
  histogram_tester.ExpectUniqueSample(
      "ChromeOS.CertProvisioning.Event.User",
      CertProvisioningEvent::kRegisteredToInvalidationTopic, 1);
  histogram_tester.ExpectTotalCount(
      "ChromeOS.CertProvisioning.KeypairGenerationTime.User", 1);
  histogram_tester.ExpectTotalCount("ChromeOS.CertProvisioning.VaTime.User", 1);
  histogram_tester.ExpectTotalCount(
      "ChromeOS.CertProvisioning.CsrSignTime.User", 1);
}

// Checks that the worker makes all necessary requests to other modules during
// success scenario when VA challenge is not received.
TEST_F(CertProvisioningWorkerTest, NoVaSuccess) {
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/false, kCertProfileRenewalPeriod);

  CertProvisioningWorkerImpl worker(
      CertScope::kUser, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
      GetResultCallback());

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_CALL(*platform_keys_service_,
                GenerateRSAKey(TokenId::kUser, kNonVaKeyModulusLengthBits,
                               /*sw_backed=*/false,
                               /*callback=*/_))
        .Times(1)
        .WillOnce(RunOnceCallback<3>(GetPublicKey(), Status::kSuccess));

    EXPECT_START_CSR_OK_WITHOUT_VA(
        ClientCertProvisioningStartCsr(kCertScopeStrUser, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        em::HashingAlgorithm::SHA256);

    EXPECT_CALL(
        *key_permissions_manager_,
        AllowKeyForUsage(/*callback=*/_, KeyUsage::kCorporate, GetPublicKey()));

    EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
        TokenId::kUser, GetPublicKey(),
        KeyAttributeType::kCertificateProvisioningId, kCertProfileId, _));

    EXPECT_SIGN_RSAPKC1_DIGEST_OK(
        SignRSAPKCS1Digest(::testing::Optional(TokenId::kUser), kDataToSign,
                           GetPublicKey(), HashAlgorithm::HASH_ALGORITHM_SHA256,
                           /*callback=*/_));

    EXPECT_FINISH_CSR_OK(ClientCertProvisioningFinishCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*va_challenge_response=*/"", kSignature, /*callback=*/_));

    EXPECT_DOWNLOAD_CERT_OK(ClientCertProvisioningDownloadCert(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));

    EXPECT_IMPORT_CERTIFICATE_OK(
        ImportCertificate(TokenId::kUser, /*certificate=*/_, /*callback=*/_));

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kSucceeded))
        .Times(1);
  }

  worker.DoStep();
}

// Checks that the worker correctly forwards a request with
// hashing_algorithm=NO_HASH to platform_keys.
TEST_F(CertProvisioningWorkerTest, NoHashInStartCsr) {
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  MockCertProvisioningInvalidator* mock_invalidator = nullptr;
  CertProvisioningWorkerImpl worker(
      CertScope::kUser, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(&mock_invalidator),
      GetStateChangeCallback(), GetResultCallback());

  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_START_CSR_OK(
        ClientCertProvisioningStartCsr(kCertScopeStrUser, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        em::HashingAlgorithm::NO_HASH);
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_CALL(*mock_invalidator, Register(kInvalidationTopic, _)).Times(1);

    EXPECT_SIGN_CHALLENGE_OK(*mock_tpm_challenge_key,
                             StartSignChallengeStep(kChallenge,
                                                    /*callback=*/_));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_REGISTER_KEY_OK(*mock_tpm_challenge_key, StartRegisterKeyStep);
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_CALL(
        *key_permissions_manager_,
        AllowKeyForUsage(/*callback=*/_, KeyUsage::kCorporate, GetPublicKey()));

    EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
        TokenId::kUser, GetPublicKey(),
        KeyAttributeType::kCertificateProvisioningId, kCertProfileId, _));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_SIGN_RSAPKC1_RAW_OK(
        SignRSAPKCS1Raw(::testing::Optional(TokenId::kUser), kDataToSign,
                        GetPublicKey(), /*callback=*/_));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_FINISH_CSR_OK(ClientCertProvisioningFinishCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        kChallengeResponse, kSignature, /*callback=*/_));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_DOWNLOAD_CERT_OK(ClientCertProvisioningDownloadCert(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));

    EXPECT_IMPORT_CERTIFICATE_OK(
        ImportCertificate(TokenId::kUser, /*certificate=*/_, /*callback=*/_));
    EXPECT_CALL(state_change_callback_observer_, StateChangeCallback());

    EXPECT_CALL(*mock_invalidator, Unregister()).Times(1);

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kSucceeded))
        .Times(1);
  }

  worker.DoStep();
}

// Checks that when the server returns try_again_later field, the worker will
// retry a request when it asked to continue the provisioning.
TEST_F(CertProvisioningWorkerTest, TryLaterManualRetry) {
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  CertProvisioningWorkerImpl worker(
      CertScope::kDevice, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
      GetResultCallback());
  const base::TimeDelta delay = base::Seconds(30);

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_DEVICE,
                            /*will_register_key=*/true,
                            /*key_name=*/GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    EXPECT_START_CSR_TRY_LATER(
        ClientCertProvisioningStartCsr(kCertScopeStrDevice, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        delay.InMilliseconds());

    worker.DoStep();
    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kKeypairGenerated);
  }

  {
    testing::InSequence seq;

    EXPECT_START_CSR_OK(
        ClientCertProvisioningStartCsr(kCertScopeStrDevice, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        em::HashingAlgorithm::SHA256);

    EXPECT_SIGN_CHALLENGE_OK(*mock_tpm_challenge_key,
                             StartSignChallengeStep(kChallenge,
                                                    /*callback=*/_));

    EXPECT_REGISTER_KEY_OK(*mock_tpm_challenge_key, StartRegisterKeyStep);

    EXPECT_CALL(
        *key_permissions_manager_,
        AllowKeyForUsage(/*callback=*/_, KeyUsage::kCorporate, GetPublicKey()));

    EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
        TokenId::kSystem, GetPublicKey(),
        KeyAttributeType::kCertificateProvisioningId, kCertProfileId, _));

    EXPECT_SIGN_RSAPKC1_DIGEST_OK(SignRSAPKCS1Digest);

    EXPECT_FINISH_CSR_TRY_LATER(
        ClientCertProvisioningFinishCsr(
            kCertScopeStrDevice, kCertProfileId, kCertProfileVersion,
            GetPublicKey(), kChallengeResponse, kSignature, /*callback=*/_),
        delay.InMilliseconds());

    worker.DoStep();
    EXPECT_EQ(worker.GetState(), CertProvisioningWorkerState::kSignCsrFinished);
  }

  {
    testing::InSequence seq;

    EXPECT_FINISH_CSR_OK(ClientCertProvisioningFinishCsr(
        kCertScopeStrDevice, kCertProfileId, kCertProfileVersion,
        GetPublicKey(), kChallengeResponse, kSignature, /*callback=*/_));

    EXPECT_DOWNLOAD_CERT_TRY_LATER(
        ClientCertProvisioningDownloadCert(kCertScopeStrDevice, kCertProfileId,
                                           kCertProfileVersion, GetPublicKey(),
                                           /*callback=*/_),
        delay.InMilliseconds());

    worker.DoStep();
    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kFinishCsrResponseReceived);
  }

  {
    testing::InSequence seq;

    EXPECT_DOWNLOAD_CERT_OK(
        ClientCertProvisioningDownloadCert(kCertScopeStrDevice, kCertProfileId,
                                           kCertProfileVersion, GetPublicKey(),
                                           /*callback=*/_));

    EXPECT_IMPORT_CERTIFICATE_OK(
        ImportCertificate(TokenId::kSystem, /*certificate=*/_, /*callback=*/_));

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kSucceeded))
        .Times(1);

    worker.DoStep();
    EXPECT_EQ(worker.GetState(), CertProvisioningWorkerState::kSucceeded);
  }
}

// Checks that when the server returns try_again_later field, the worker will
// automatically retry a request after some time.
TEST_F(CertProvisioningWorkerTest, TryLaterWait) {
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  CertProvisioningWorkerImpl worker(
      CertScope::kUser, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
      GetResultCallback());

  const base::TimeDelta start_csr_delay = base::Seconds(30);
  const base::TimeDelta finish_csr_delay = base::Seconds(30);
  const base::TimeDelta download_cert_server_delay = base::Milliseconds(100);
  const base::TimeDelta download_cert_real_delay = base::Seconds(10);
  const base::TimeDelta small_delay = base::Milliseconds(500);

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    EXPECT_START_CSR_TRY_LATER(
        ClientCertProvisioningStartCsr(kCertScopeStrUser, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        start_csr_delay.InMilliseconds());

    worker.DoStep();
    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kKeypairGenerated);
  }

  {
    testing::InSequence seq;

    EXPECT_START_CSR_OK(
        ClientCertProvisioningStartCsr(kCertScopeStrUser, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        em::HashingAlgorithm::SHA256);

    EXPECT_SIGN_CHALLENGE_OK(*mock_tpm_challenge_key,
                             StartSignChallengeStep(kChallenge,
                                                    /*callback=*/_));

    EXPECT_REGISTER_KEY_OK(*mock_tpm_challenge_key, StartRegisterKeyStep);

    EXPECT_CALL(
        *key_permissions_manager_,
        AllowKeyForUsage(/*callback=*/_, KeyUsage::kCorporate, GetPublicKey()));

    EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
        TokenId::kUser, GetPublicKey(),
        KeyAttributeType::kCertificateProvisioningId, kCertProfileId, _));

    EXPECT_SIGN_RSAPKC1_DIGEST_OK(
        SignRSAPKCS1Digest(::testing::Optional(TokenId::kUser), kDataToSign,
                           GetPublicKey(), HashAlgorithm::HASH_ALGORITHM_SHA256,
                           /*callback=*/_));

    EXPECT_FINISH_CSR_TRY_LATER(
        ClientCertProvisioningFinishCsr(
            kCertScopeStrUser, kCertProfileId, kCertProfileVersion,
            GetPublicKey(), kChallengeResponse, kSignature, /*callback=*/_),
        finish_csr_delay.InMilliseconds());

    FastForwardBy(start_csr_delay + small_delay);
    EXPECT_EQ(worker.GetState(), CertProvisioningWorkerState::kSignCsrFinished);
  }

  {
    testing::InSequence seq;

    EXPECT_FINISH_CSR_OK(ClientCertProvisioningFinishCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        kChallengeResponse, kSignature, /*callback=*/_));

    EXPECT_DOWNLOAD_CERT_TRY_LATER(
        ClientCertProvisioningDownloadCert(kCertScopeStrUser, kCertProfileId,
                                           kCertProfileVersion, GetPublicKey(),
                                           /*callback=*/_),
        download_cert_server_delay.InMilliseconds());

    FastForwardBy(finish_csr_delay + small_delay);
    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kFinishCsrResponseReceived);
  }

  {
    testing::InSequence seq;

    EXPECT_DOWNLOAD_CERT_OK(ClientCertProvisioningDownloadCert);

    EXPECT_IMPORT_CERTIFICATE_OK(
        ImportCertificate(TokenId::kUser, /*certificate=*/_, /*callback=*/_));

    FastForwardBy(small_delay);
    // Check that minimum wait time is not too small even if the server
    // has responded with a small one.
    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kFinishCsrResponseReceived);

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kSucceeded))
        .Times(1);
    FastForwardBy(download_cert_real_delay + small_delay);
    EXPECT_EQ(worker.GetState(), CertProvisioningWorkerState::kSucceeded);
  }
}

// Checks that when the device management server returns a
// DM_STATUS_SERVICE_ACTIVATION_PENDING status error (which is 412 pending
// approval) the server retries the request after the expected delay depending
// on the request.
TEST_F(CertProvisioningWorkerTest, ServiceActivationPendingResponse) {
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  CertProvisioningWorkerImpl worker(
      CertScope::kUser, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
      GetResultCallback());

  const base::TimeDelta kSmallDelay = base::Milliseconds(500);
  const base::TimeDelta kExpectedStartCsrDelay = base::Hours(1);
  const base::TimeDelta kExpectedFinishCsrDelay = base::Hours(1);
  const base::TimeDelta kExpectedDownloadCsrDelay = base::Hours(8);

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    EXPECT_START_CSR_SERVICE_ACTIVATION_PENDING(ClientCertProvisioningStartCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));

    worker.DoStep();
    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kKeypairGenerated);
  }

  {
    testing::InSequence seq;

    // Verify that nothing happens after half of the expected StartCsr delay.
    FastForwardBy(kExpectedStartCsrDelay / 2);
    Mock::VerifyAndClearExpectations(&cloud_policy_client_);

    EXPECT_START_CSR_OK(
        ClientCertProvisioningStartCsr(kCertScopeStrUser, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        em::HashingAlgorithm::SHA256);

    EXPECT_SIGN_CHALLENGE_OK(*mock_tpm_challenge_key,
                             StartSignChallengeStep(kChallenge,
                                                    /*callback=*/_));

    EXPECT_REGISTER_KEY_OK(*mock_tpm_challenge_key, StartRegisterKeyStep);

    EXPECT_CALL(
        *key_permissions_manager_,
        AllowKeyForUsage(/*callback=*/_, KeyUsage::kCorporate, GetPublicKey()));

    EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
        TokenId::kUser, GetPublicKey(),
        KeyAttributeType::kCertificateProvisioningId, kCertProfileId, _));

    EXPECT_SIGN_RSAPKC1_DIGEST_OK(
        SignRSAPKCS1Digest(::testing::Optional(TokenId::kUser), kDataToSign,
                           GetPublicKey(), HashAlgorithm::HASH_ALGORITHM_SHA256,
                           /*callback=*/_));

    EXPECT_FINISH_CSR_SERVICE_ACTIVATION_PENDING(
        ClientCertProvisioningFinishCsr(
            kCertScopeStrUser, kCertProfileId, kCertProfileVersion,
            GetPublicKey(), kChallengeResponse, kSignature, /*callback=*/_));

    FastForwardBy(kExpectedStartCsrDelay / 2 + kSmallDelay);
    EXPECT_EQ(worker.GetState(), CertProvisioningWorkerState::kSignCsrFinished);
  }

  {
    testing::InSequence seq;

    // Verify that nothing happens after half of the expected FinishCsr delay.
    FastForwardBy(kExpectedFinishCsrDelay / 2);
    Mock::VerifyAndClearExpectations(&cloud_policy_client_);

    EXPECT_FINISH_CSR_OK(ClientCertProvisioningFinishCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        kChallengeResponse, kSignature, /*callback=*/_));

    EXPECT_DOWNLOAD_CERT_SERVICE_ACTIVATION_PENDING(
        ClientCertProvisioningDownloadCert(kCertScopeStrUser, kCertProfileId,
                                           kCertProfileVersion, GetPublicKey(),
                                           /*callback=*/_));

    FastForwardBy(kExpectedFinishCsrDelay / 2 + kSmallDelay);
    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kFinishCsrResponseReceived);
  }

  {
    testing::InSequence seq;

    // Verify that nothing happens after half of the expected DownloadCert
    // delay.
    FastForwardBy(kExpectedDownloadCsrDelay / 2);
    Mock::VerifyAndClearExpectations(&cloud_policy_client_);

    EXPECT_DOWNLOAD_CERT_OK(ClientCertProvisioningDownloadCert);

    EXPECT_IMPORT_CERTIFICATE_OK(
        ImportCertificate(TokenId::kUser, /*certificate=*/_, /*callback=*/_));

    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kFinishCsrResponseReceived);

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kSucceeded))
        .Times(1);
    FastForwardBy(kExpectedDownloadCsrDelay / 2 + kSmallDelay);
    EXPECT_EQ(worker.GetState(), CertProvisioningWorkerState::kSucceeded);
  }
}

// Checks that when the server returns try_again_later field, the worker will
// retry when the invalidation is triggered.
TEST_F(CertProvisioningWorkerTest, InvalidationRespected) {
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  MockCertProvisioningInvalidator* mock_invalidator = nullptr;
  CertProvisioningWorkerImpl worker(
      CertScope::kUser, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(&mock_invalidator),
      GetStateChangeCallback(), GetResultCallback());

  const base::TimeDelta start_csr_delay = base::Seconds(30);
  const base::TimeDelta finish_csr_delay = base::Seconds(30);
  const base::TimeDelta download_cert_server_delay = base::Milliseconds(100);
  const base::TimeDelta small_delay = base::Milliseconds(500);

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    EXPECT_START_CSR_TRY_LATER(
        ClientCertProvisioningStartCsr(kCertScopeStrUser, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        start_csr_delay.InMilliseconds());

    worker.DoStep();
    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kKeypairGenerated);
  }

  base::RepeatingClosure on_invalidation_callback;
  {
    testing::InSequence seq;

    EXPECT_START_CSR_OK(
        ClientCertProvisioningStartCsr(kCertScopeStrUser, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        em::HashingAlgorithm::SHA256);
    EXPECT_CALL(*mock_invalidator, Register(kInvalidationTopic, _))
        .WillOnce(SaveArg<1>(&on_invalidation_callback));

    EXPECT_SIGN_CHALLENGE_OK(*mock_tpm_challenge_key,
                             StartSignChallengeStep(kChallenge,
                                                    /*callback=*/_));

    EXPECT_REGISTER_KEY_OK(*mock_tpm_challenge_key, StartRegisterKeyStep);

    EXPECT_CALL(
        *key_permissions_manager_,
        AllowKeyForUsage(/*callback=*/_, KeyUsage::kCorporate, GetPublicKey()));

    EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
        TokenId::kUser, GetPublicKey(),
        KeyAttributeType::kCertificateProvisioningId, kCertProfileId, _));

    EXPECT_SIGN_RSAPKC1_DIGEST_OK(
        SignRSAPKCS1Digest(::testing::Optional(TokenId::kUser), kDataToSign,
                           GetPublicKey(), HashAlgorithm::HASH_ALGORITHM_SHA256,
                           /*callback=*/_));

    EXPECT_FINISH_CSR_TRY_LATER(
        ClientCertProvisioningFinishCsr(
            kCertScopeStrUser, kCertProfileId, kCertProfileVersion,
            GetPublicKey(), kChallengeResponse, kSignature, /*callback=*/_),
        finish_csr_delay.InMilliseconds());

    FastForwardBy(start_csr_delay + small_delay);
    EXPECT_EQ(worker.GetState(), CertProvisioningWorkerState::kSignCsrFinished);
  }

  {
    testing::InSequence seq;

    EXPECT_FINISH_CSR_OK(ClientCertProvisioningFinishCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        kChallengeResponse, kSignature, /*callback=*/_));

    EXPECT_DOWNLOAD_CERT_TRY_LATER(
        ClientCertProvisioningDownloadCert(kCertScopeStrUser, kCertProfileId,
                                           kCertProfileVersion, GetPublicKey(),
                                           /*callback=*/_),
        download_cert_server_delay.InMilliseconds());

    FastForwardBy(finish_csr_delay + small_delay);
    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kFinishCsrResponseReceived);
  }

  {
    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kFinishCsrResponseReceived);

    testing::InSequence seq;

    EXPECT_DOWNLOAD_CERT_OK(ClientCertProvisioningDownloadCert);

    EXPECT_IMPORT_CERTIFICATE_OK(
        ImportCertificate(TokenId::kUser, /*certificate=*/_, /*callback=*/_));

    EXPECT_CALL(*mock_invalidator, Unregister()).Times(1);

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kSucceeded))
        .Times(1);

    on_invalidation_callback.Run();
    EXPECT_EQ(worker.GetState(), CertProvisioningWorkerState::kSucceeded);
  }
}

// Checks that when the server returns error status, the worker will enter an
// error state and stop the provisioning.
TEST_F(CertProvisioningWorkerTest, StatusErrorHandling) {
  const CertScope kCertScope = CertScope::kUser;
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  CertProvisioningWorkerImpl worker(
      kCertScope, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
      GetResultCallback());

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    EXPECT_START_CSR_INVALID_REQUEST(ClientCertProvisioningStartCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kFailed))
        .Times(1);
  }

  worker.DoStep();
  FastForwardBy(base::Seconds(1));

  VerifyDeleteKeyCalledOnce(kCertScope);
}

// Checks that when the server returns response error, the worker will enter an
// error state and stop the provisioning. Also check factory.
TEST_F(CertProvisioningWorkerTest, ResponseErrorHandling) {
  const CertScope kCertScope = CertScope::kUser;
  base::HistogramTester histogram_tester;

  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  auto worker = CertProvisioningWorkerFactory::Get()->Create(
      kCertScope, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
      GetResultCallback());

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    EXPECT_START_CSR_CA_ERROR(ClientCertProvisioningStartCsr);

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kFailed))
        .Times(1);
  }

  worker->DoStep();
  FastForwardBy(base::Seconds(1));

  VerifyDeleteKeyCalledOnce(kCertScope);

  histogram_tester.ExpectBucketCount("ChromeOS.CertProvisioning.Result.User",
                                     CertProvisioningWorkerState::kFailed, 1);
  histogram_tester.ExpectBucketCount(
      "ChromeOS.CertProvisioning.Result.User",
      CertProvisioningWorkerState::kKeypairGenerated, 1);
  histogram_tester.ExpectTotalCount("ChromeOS.CertProvisioning.Result.User", 2);
}

TEST_F(CertProvisioningWorkerTest, InconsistentDataErrorHandling) {
  const CertScope kCertScope = CertScope::kUser;
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  auto worker = CertProvisioningWorkerFactory::Get()->Create(
      kCertScope, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
      GetResultCallback());

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    EXPECT_START_CSR_INCONSISTENT_DATA(ClientCertProvisioningStartCsr);

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile,
                         CertProvisioningWorkerState::kInconsistentDataError))
        .Times(1);
  }

  worker->DoStep();
  FastForwardBy(base::Seconds(1));

  VerifyDeleteKeyCalledOnce(kCertScope);
}

// Checks that when the server returns TEMPORARY_UNAVAILABLE status code, the
// worker will automatically retry a request using exponential backoff strategy.
TEST_F(CertProvisioningWorkerTest, BackoffStrategy) {
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  CertProvisioningWorkerImpl worker(
      CertScope::kUser, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
      GetResultCallback());

  base::TimeDelta next_delay = base::Seconds(30);
  const base::TimeDelta small_delay = base::Milliseconds(500);

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    EXPECT_START_CSR_TEMPORARY_UNAVAILABLE(ClientCertProvisioningStartCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));
    worker.DoStep();
  }

  Mock::VerifyAndClearExpectations(&cloud_policy_client_);

  {
    EXPECT_START_CSR_TEMPORARY_UNAVAILABLE(ClientCertProvisioningStartCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));
    FastForwardBy(next_delay + small_delay * 10);
    next_delay *= 2;
  }

  Mock::VerifyAndClearExpectations(&cloud_policy_client_);

  {
    EXPECT_START_CSR_TEMPORARY_UNAVAILABLE(ClientCertProvisioningStartCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));
    FastForwardBy(next_delay + small_delay * 10);
    next_delay *= 2;
  }

  Mock::VerifyAndClearExpectations(&cloud_policy_client_);

  {
    EXPECT_START_CSR_TEMPORARY_UNAVAILABLE(ClientCertProvisioningStartCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));
    FastForwardBy(next_delay + small_delay);
    next_delay *= 2;
  }
}

// Checks that the worker removes a key when an error occurs after the key was
// registered.
TEST_F(CertProvisioningWorkerTest, RemoveRegisteredKey) {
  base::HistogramTester histogram_tester;

  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);
  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  MockCertProvisioningInvalidator* mock_invalidator = nullptr;
  CertProvisioningWorkerImpl worker(
      CertScope::kUser, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(&mock_invalidator),
      GetStateChangeCallback(), GetResultCallback());

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    EXPECT_START_CSR_OK(
        ClientCertProvisioningStartCsr(kCertScopeStrUser, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        em::HashingAlgorithm::SHA256);

    EXPECT_CALL(*mock_invalidator, Register(kInvalidationTopic, _)).Times(1);

    EXPECT_SIGN_CHALLENGE_OK(*mock_tpm_challenge_key,
                             StartSignChallengeStep(kChallenge,
                                                    /*callback=*/_));

    EXPECT_REGISTER_KEY_OK(*mock_tpm_challenge_key, StartRegisterKeyStep);

    EXPECT_CALL(
        *key_permissions_manager_,
        AllowKeyForUsage(/*callback=*/_, KeyUsage::kCorporate, GetPublicKey()));

    EXPECT_SET_ATTRIBUTE_FOR_KEY_FAIL(SetAttributeForKey(
        TokenId::kUser, GetPublicKey(),
        KeyAttributeType::kCertificateProvisioningId, kCertProfileId, _));

    EXPECT_CALL(*mock_invalidator, Unregister()).Times(1);

    EXPECT_CALL(
        *platform_keys_service_,
        RemoveKey(TokenId::kUser,
                  /*public_key_spki_der=*/GetPublicKey(), /*callback=*/_))
        .Times(1)
        .WillOnce(RunOnceCallback<2>(Status::kSuccess));

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kFailed))
        .Times(1);
  }

  worker.DoStep();
  FastForwardBy(base::Seconds(1));

  histogram_tester.ExpectBucketCount("ChromeOS.CertProvisioning.Result.User",
                                     CertProvisioningWorkerState::kFailed, 1);
  histogram_tester.ExpectBucketCount(
      "ChromeOS.CertProvisioning.Result.User",
      CertProvisioningWorkerState::kKeyRegistered, 1);
  histogram_tester.ExpectTotalCount("ChromeOS.CertProvisioning.Result.User", 2);
}

class PrefServiceObserver {
 public:
  PrefServiceObserver(PrefService* service, const char* pref_name)
      : service_(service), pref_name_(pref_name) {
    pref_change_registrar_.Init(service);
    pref_change_registrar_.Add(
        pref_name, base::BindRepeating(&PrefServiceObserver::OnPrefsChange,
                                       weak_factory_.GetWeakPtr()));
  }

  void OnPrefsChange() {
    const base::Value* pref_value = service_->Get(pref_name_);
    DCHECK(pref_value);
    OnPrefValueUpdated(*pref_value);
  }

  // Allows to add expectations about preference changes and verify new values.
  MOCK_METHOD(void, OnPrefValueUpdated, (const base::Value& value));

 private:
  PrefService* service_ = nullptr;
  const char* pref_name_ = nullptr;
  PrefChangeRegistrar pref_change_registrar_;
  base::WeakPtrFactory<PrefServiceObserver> weak_factory_{this};
};

TEST_F(CertProvisioningWorkerTest, SerializationSuccess) {
  const base::TimeDelta kRenewalPeriod = base::Seconds(1200300);
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kRenewalPeriod);
  const CertScope kCertScope = CertScope::kUser;

  std::unique_ptr<MockCertProvisioningInvalidator> mock_invalidator_obj;
  MockCertProvisioningInvalidator* mock_invalidator = nullptr;

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  std::unique_ptr<CertProvisioningWorker> worker =
      CertProvisioningWorkerFactory::Get()->Create(
          kCertScope, GetProfile(), &testing_pref_service_, cert_profile,
          &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
          GetResultCallback());

  StrictMock<PrefServiceObserver> pref_observer(
      &testing_pref_service_, GetPrefNameForSerialization(kCertScope));
  base::Value pref_val;

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));

  // Prepare key, send start csr request.
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    pref_val = ParseJson(base::StringPrintf(
        R"({
          "cert_profile_1": {
            "cert_profile": {
              "policy_version": "cert_profile_version_1",
              "name": "Certificate Profile 1",
              "profile_id": "cert_profile_1",
              "va_enabled": true,
              "renewal_period": 1200300
            },
            "cert_scope": 0,
            "invalidation_topic": "",
            "public_key": "%s",
            "state": 1
          }
        })",
        kPublicKeyBase64));
    EXPECT_CALL(pref_observer, OnPrefValueUpdated(IsJson(pref_val))).Times(1);

    EXPECT_START_CSR_NO_OP(ClientCertProvisioningStartCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));

    worker->DoStep();
  }

  // Recreate worker.
  {
    testing::InSequence seq;

    mock_tpm_challenge_key = PrepareTpmChallengeKey();

    EXPECT_CALL(*mock_tpm_challenge_key,
                RestorePreparedKeyState(
                    attestation::AttestationKeyType::KEY_USER,
                    /*will_register_key=*/true, GetKeyName(kCertProfileId),
                    GetPublicKey(), /*profile=*/_))
        .Times(1);

    worker = CertProvisioningWorkerFactory::Get()->Deserialize(
        kCertScope, GetProfile(), &testing_pref_service_,
        *pref_val.FindKeyOfType(kCertProfileId, base::Value::Type::DICTIONARY),
        &cloud_policy_client_, MakeInvalidator(&mock_invalidator),
        GetStateChangeCallback(), GetResultCallback());
  }

  // Retry start csr request, receive response, try sign challenge.
  {
    testing::InSequence seq;

    EXPECT_START_CSR_OK(
        ClientCertProvisioningStartCsr(kCertScopeStrUser, kCertProfileId,
                                       kCertProfileVersion, GetPublicKey(),
                                       /*callback=*/_),
        em::HashingAlgorithm::SHA256);

    pref_val = ParseJson("{}");
    EXPECT_CALL(pref_observer, OnPrefValueUpdated(IsJson(pref_val))).Times(1);

    EXPECT_CALL(*mock_invalidator, Register(kInvalidationTopic, _)).Times(1);

    EXPECT_SIGN_CHALLENGE_OK(*mock_tpm_challenge_key,
                             StartSignChallengeStep(kChallenge,
                                                    /*callback=*/_));

    EXPECT_REGISTER_KEY_OK(*mock_tpm_challenge_key, StartRegisterKeyStep);

    EXPECT_CALL(
        *key_permissions_manager_,
        AllowKeyForUsage(/*callback=*/_, KeyUsage::kCorporate, GetPublicKey()));

    EXPECT_SET_ATTRIBUTE_FOR_KEY_OK(SetAttributeForKey(
        TokenId::kUser, GetPublicKey(),
        KeyAttributeType::kCertificateProvisioningId, kCertProfileId, _));

    EXPECT_SIGN_RSAPKC1_DIGEST_OK(
        SignRSAPKCS1Digest(::testing::Optional(TokenId::kUser), kDataToSign,
                           GetPublicKey(), HashAlgorithm::HASH_ALGORITHM_SHA256,
                           /*callback=*/_));

    EXPECT_FINISH_CSR_OK(ClientCertProvisioningFinishCsr(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        kChallengeResponse, kSignature, /*callback=*/_));

    pref_val = ParseJson(base::StringPrintf(
        R"({
          "cert_profile_1": {
            "cert_profile": {
              "policy_version": "cert_profile_version_1",
              "name": "Certificate Profile 1",
              "profile_id": "cert_profile_1",
              "va_enabled": true,
              "renewal_period": 1200300
            },
            "cert_scope": 0,
            "invalidation_topic": "fake_invalidation_topic_1",
            "public_key": "%s",
            "state": 7
          }
        })",
        kPublicKeyBase64));
    EXPECT_CALL(pref_observer, OnPrefValueUpdated(IsJson(pref_val))).Times(1);

    EXPECT_DOWNLOAD_CERT_NO_OP(ClientCertProvisioningDownloadCert(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));

    worker->DoStep();
  }

  // Recreate worker.
  {
    testing::InSequence seq;

    mock_invalidator_obj = MakeInvalidator(&mock_invalidator);
    EXPECT_CALL(*mock_invalidator, Register(kInvalidationTopic, _)).Times(1);

    mock_tpm_challenge_key = PrepareTpmChallengeKey();
    EXPECT_CALL(*mock_tpm_challenge_key,
                RestorePreparedKeyState(
                    attestation::AttestationKeyType::KEY_USER,
                    /*will_register_key=*/true, GetKeyName(kCertProfileId),
                    GetPublicKey(), /*profile=*/_))
        .Times(1);

    worker = CertProvisioningWorkerFactory::Get()->Deserialize(
        kCertScope, GetProfile(), &testing_pref_service_,
        *pref_val.FindKeyOfType(kCertProfileId, base::Value::Type::DICTIONARY),
        &cloud_policy_client_, std::move(mock_invalidator_obj),
        GetStateChangeCallback(), GetResultCallback());
  }

  // Retry download cert request, receive response, try import certificate.
  {
    testing::InSequence seq;

    EXPECT_DOWNLOAD_CERT_OK(ClientCertProvisioningDownloadCert(
        kCertScopeStrUser, kCertProfileId, kCertProfileVersion, GetPublicKey(),
        /*callback=*/_));

    EXPECT_IMPORT_CERTIFICATE_OK(
        ImportCertificate(TokenId::kUser, /*certificate=*/_, /*callback=*/_));

    pref_val = ParseJson("{}");
    EXPECT_CALL(pref_observer, OnPrefValueUpdated(IsJson(pref_val))).Times(1);

    EXPECT_CALL(*mock_invalidator, Unregister()).Times(1);

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kSucceeded))
        .Times(1);
    worker->DoStep();
  }
}

TEST_F(CertProvisioningWorkerTest, SerializationOnFailure) {
  const CertScope kCertScope = CertScope::kUser;
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  auto worker = CertProvisioningWorkerFactory::Get()->Create(
      kCertScope, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
      GetResultCallback());

  PrefServiceObserver pref_observer(&testing_pref_service_,
                                    GetPrefNameForSerialization(kCertScope));
  base::Value pref_val;

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_USER,
                            /*will_register_key=*/true,
                            GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    pref_val = ParseJson(base::StringPrintf(
        R"({
          "cert_profile_1": {
            "cert_profile": {
              "policy_version": "cert_profile_version_1",
              "name": "Certificate Profile 1",
              "profile_id": "cert_profile_1",
              "va_enabled": true
            },
            "cert_scope": 0,
            "invalidation_topic": "",
            "public_key": "%s",
            "state": 1
          }
        })",
        kPublicKeyBase64));
    EXPECT_CALL(pref_observer, OnPrefValueUpdated(IsJson(pref_val))).Times(1);

    EXPECT_START_CSR_CA_ERROR(ClientCertProvisioningStartCsr);

    pref_val = ParseJson("{}");
    EXPECT_CALL(pref_observer, OnPrefValueUpdated(IsJson(pref_val))).Times(1);

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kFailed))
        .Times(1);
  }

  worker->DoStep();
  FastForwardBy(base::Seconds(1));

  VerifyDeleteKeyCalledOnce(kCertScope);
}

TEST_F(CertProvisioningWorkerTest, InformationalGetters) {
  const CertScope kCertScope = CertScope::kUser;
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  CertProvisioningWorkerImpl worker(
      kCertScope, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
      GetResultCallback());

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(*mock_tpm_challenge_key, StartPrepareKeyStep);

    EXPECT_START_CSR_TRY_LATER(ClientCertProvisioningStartCsr,
                               base::Seconds(30).InMilliseconds());

    worker.DoStep();
    EXPECT_EQ(worker.GetState(),
              CertProvisioningWorkerState::kKeypairGenerated);
    EXPECT_EQ(worker.GetPreviousState(),
              CertProvisioningWorkerState::kInitState);
    EXPECT_EQ(worker.GetCertProfile(), cert_profile);
    EXPECT_EQ(worker.GetPublicKey(), GetPublicKey());
  }

  {
    testing::InSequence seq;

    EXPECT_START_CSR_CA_ERROR(ClientCertProvisioningStartCsr);

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kFailed))
        .Times(1);

    worker.DoStep();
    FastForwardBy(base::Seconds(1));

    VerifyDeleteKeyCalledOnce(kCertScope);

    EXPECT_EQ(worker.GetState(), CertProvisioningWorkerState::kFailed);
    EXPECT_EQ(worker.GetPreviousState(),
              CertProvisioningWorkerState::kKeypairGenerated);
    EXPECT_EQ(worker.GetCertProfile(), cert_profile);
    EXPECT_EQ(worker.GetPublicKey(), GetPublicKey());
  }
}

TEST_F(CertProvisioningWorkerTest, CancelDeviceWorker) {
  base::HistogramTester histogram_tester;

  const CertScope kCertScope = CertScope::kDevice;
  CertProfile cert_profile(kCertProfileId, kCertProfileName,
                           kCertProfileVersion,
                           /*is_va_enabled=*/true, kCertProfileRenewalPeriod);

  EXPECT_CALL(state_change_callback_observer_, StateChangeCallback)
      .Times(AtLeast(1));
  MockTpmChallengeKeySubtle* mock_tpm_challenge_key = PrepareTpmChallengeKey();
  auto worker = CertProvisioningWorkerFactory::Get()->Create(
      kCertScope, GetProfile(), &testing_pref_service_, cert_profile,
      &cloud_policy_client_, MakeInvalidator(), GetStateChangeCallback(),
      GetResultCallback());

  EXPECT_CALL(callback_observer_, Callback).Times(0);

  PrefServiceObserver pref_observer(&testing_pref_service_,
                                    GetPrefNameForSerialization(kCertScope));
  base::Value pref_val;

  {
    testing::InSequence seq;

    EXPECT_PREPARE_KEY_OK(
        *mock_tpm_challenge_key,
        StartPrepareKeyStep(attestation::AttestationKeyType::KEY_DEVICE,
                            /*will_register_key=*/true,
                            /*key_name=*/GetKeyName(kCertProfileId),
                            /*profile=*/_,
                            /*callback=*/_, /*signals=*/_));

    pref_val = ParseJson(base::StringPrintf(
        R"({
          "cert_profile_1": {
            "cert_profile": {
              "policy_version": "cert_profile_version_1",
              "name": "Certificate Profile 1",
              "profile_id": "cert_profile_1",
              "va_enabled": true
            },
            "cert_scope": 1,
            "invalidation_topic": "",
            "public_key": "%s",
            "state": 1
          }
        })",
        kPublicKeyBase64));
    EXPECT_CALL(pref_observer, OnPrefValueUpdated(IsJson(pref_val))).Times(1);

    EXPECT_START_CSR_NO_OP(ClientCertProvisioningStartCsr);

    worker->DoStep();
  }

  {
    pref_val = ParseJson("{}");
    EXPECT_CALL(pref_observer, OnPrefValueUpdated(IsJson(pref_val))).Times(1);

    worker->Stop(CertProvisioningWorkerState::kCanceled);

    EXPECT_CALL(callback_observer_,
                Callback(cert_profile, CertProvisioningWorkerState::kCanceled))
        .Times(1);
    FastForwardBy(base::Seconds(1));

    VerifyDeleteKeyCalledOnce(kCertScope);
  }

  histogram_tester.ExpectUniqueSample("ChromeOS.CertProvisioning.Result.Device",
                                      CertProvisioningWorkerState::kCanceled,
                                      1);
}

}  // namespace
}  // namespace cert_provisioning
}  // namespace ash
